{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizing Existing Networks\n",
    "\n",
    "Nengo DL is not confined to opimizing cutom made networks, it can also be used to make existing networks better, or achieve the same result with fewer neurons. What this example will show is how to train a circular convolution network.\n",
    "\n",
    "Circular convolution is a key operation used to process [semantic pointers](http://compneuro.uwaterloo.ca/research/spa/semantic-pointer-architecture.html). By optimizing this smaller network, larger more complex networks that utilize circular convolution can benefit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import nengo\n",
    "import nengo_dl\n",
    "%load_ext nengo.ipynb\n",
    "from nengo.spa import Vocabulary\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To properly train the network, we generate novel training data by randomly generating semantic pointers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def gen_pointers(n_inputs, dims, seed):\n",
    "    state = np.random.RandomState(seed)\n",
    "    vocabulary = Vocabulary(dimensions=dims, rng=state, max_similarity=1)\n",
    "    pairs = []\n",
    "    for v in range(n_inputs):\n",
    "        # keys start with A, second element starts with B, third starts with C\n",
    "        conv_key = 'C' + str(v)\n",
    "        point_key_1= 'A' + str(v)\n",
    "        pointer_1 = vocabulary.create_pointer()\n",
    "        point_key_2 = 'B' + str(v)\n",
    "        pointer_2 = vocabulary.create_pointer()\n",
    "        vocabulary.add(point_key_1, pointer_1)\n",
    "        vocabulary.add(point_key_2, pointer_2)\n",
    "        vocabulary.add(conv_key, vocabulary.parse(point_key_2 + \"*\" + point_key_1))\n",
    "    \n",
    "    A = np.asarray([vocabulary['A' + str(i)].v for i in range(n_inputs)])[:, None, :]\n",
    "    B = np.asarray([vocabulary['B' + str(i)].v for i in range(n_inputs)])[:, None, :]\n",
    "    C = np.asarray([vocabulary['C' + str(i)].v for i in range(n_inputs)])[:, None, :]\n",
    "    return A, B, C, vocabulary\n",
    "\n",
    "\n",
    "dimensions = 50\n",
    "seed = 0\n",
    "\n",
    "test_a, test_b, test_c, vocab = gen_pointers(10, dimensions, seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want our optimized network to work with spiking LIF neurons, so we will use SoftLIFRate neurons (a differentiable approximation of LIF neurons) to train the network.\n",
    "\n",
    "We'll start with the `nengo.networks.CircularConvolution` network, where all the parameters are initialized using the standard Nengo methods, and then further optimize those parameters using deep learning training methods.\n",
    "\n",
    "In this example only 5 neurons are used per dimension for the circular convolution. This is fewer than would typically be used in a Nengo model, but the enhanced performance enabled by the training process will allow the network to function well with this restricted number of neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with nengo.Network(seed=seed) as net:\n",
    "    net.config[nengo.Ensemble].neuron_type = nengo_dl.SoftLIFRate(sigma=0.1)\n",
    "    net.config[nengo.Connection].synapse = None\n",
    "\n",
    "    # Get the raw vectors for the pointers using `vocab['A'].v`\n",
    "    a = nengo.Node(output=vocab['A0'].v)\n",
    "    b = nengo.Node(output=vocab['B0'].v)\n",
    "\n",
    "    # Make the circular convolution network with 5 neurons per dimension\n",
    "    cconv = nengo.networks.CircularConvolution(5, dimensions=dimensions)\n",
    "\n",
    "    # Connect the input nodes to the input slots `A` and `B` on the network\n",
    "    nengo.Connection(a, cconv.input_a)\n",
    "    nengo.Connection(b, cconv.input_b)\n",
    "\n",
    "    # Probe the output\n",
    "    out = nengo.Probe(cconv.output)\n",
    "    out_filtered = nengo.Probe(cconv.output, synapse=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now run the network in its default state to get an idea of the baseline performance. Ideally the output would be clearly `C0`, the result of the convolution between `A0` and `B0`, but we can see that it is poorly differentiated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(net) as sim:\n",
    "    sim.run(0.3)\n",
    "plt.figure()\n",
    "output_vocab = vocab.create_subset([\"C%d\" % i for i in range(10)])\n",
    "plt.plot(sim.trange(), nengo.spa.similarity(sim.data[out_filtered], output_vocab))\n",
    "plt.legend(output_vocab.keys, loc=4)\n",
    "plt.ylim([-1, 1])\n",
    "plt.xlabel(\"t [s]\")\n",
    "plt.ylabel(\"dot product\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can optimize our network, by showing it random input pointers and training it to output their circular convolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with nengo_dl.Simulator(net, minibatch_size=100, device=\"/cpu:0\") as sim:\n",
    "    optimizer = tf.train.RMSPropOptimizer(5e-3)\n",
    "        \n",
    "    # generate random data\n",
    "    train_a, train_b, train_c, _ = gen_pointers(1000, dimensions, seed+1)\n",
    "    input_feed = {a: train_a, b: train_b}\n",
    "    output_feed = {out: train_c}\n",
    "\n",
    "    # train the network for one epoch\n",
    "    sim.train(input_feed, output_feed, optimizer, n_epochs=100, objective=\"mse\")\n",
    "\n",
    "    sim.run(0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "After training we run the same test on the network and plot the output. Now we can clearly see that the output of the network is closest to the ideal output, `C0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = sim.data[out_filtered]\n",
    "plt.figure()\n",
    "plt.plot(sim.trange(), nengo.spa.similarity(output[0], output_vocab))\n",
    "plt.legend(output_vocab.keys, loc=4)\n",
    "plt.ylim([-1, 1])\n",
    "plt.xlabel(\"t [s]\")\n",
    "plt.ylabel(\"dot product\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a future example we will show how to integrate these training improvments into a larger network and improve the performance of the network as a whole."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
